# ------------------------------------------------------------------------------
# Copyright (c) HQU
# Licensed under the HQU License.
# Written by Wang Youjije (youjieWang@stu.hqu.edu.cn)
# Modified by Wang Youjije (youjieWang@stu.hqu.edu.cn)
# ------------------------------------------------------------------------------


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import time
import logging
import os

import numpy as np
import torch

# 在train.py中添加了init_path就不需要加lib.这个前缀了
from lib.core.evaluate import accuracy

logger = logging.getLogger(__name__)


def train(config, train_loader, model, criterions: list, optimizer, epoch,
          output_dir, tb_log_dir, writer_dict):
    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    acc = AverageMeter()

    # switch to train mode
    model.train()

    end = time.time()

    for i, (input, target_h, target_c, target_weight, meta) in enumerate(train_loader):
        # 这边输入一个， 输出两个，所以有两个目标值target_h：热图的目标值， target_c坐标的目标值
        # measure data loading time
        data_time.update(time.time() - end)

        # compute output
        outputs_h, outputs_c = model(input)  # 数据输入到网络里面,这里的输出应该是两个，一个是热图的输出，一个是坐标的输出

        target_h = target_h.cuda(non_blocking=True)
        target_c = target_c.cuda(non_blocking=True)  # TODO 4/26坐标的目标值要怎么设计，还在考虑cuda的参数不了解non_blocking=True
        target_weight = target_weight.cuda(non_blocking=True)

        if isinstance(outputs_h, list):
            loss_h = criterions[0](outputs_h[0], target_h, target_weight)
            for output_h in outputs_h[1:]:
                loss_h += criterions[0](output_h, target_h, target_weight)
        else:
            output_h = outputs_h
            loss_h = criterions[0](output_h, target_h, target_weight)

        if isinstance(outputs_c, list):
            loss_c = criterions[1](outputs_c[0], target, target_weight)
            for output_c in outputs_c[1:]:
                loss_c += criterions[1](output_c, target, target_weight)
        else:
            output_c = outputs_c
            loss_c = criterions[1](output_c, target, target_weight)

        loss = 4 * loss_h + loss_c  # TODO 这边loss的权重比不知道怎么设置

        # compute gradient and do update step
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # measure accuracy and record loss
        losses.update(loss.item(), input.size(0))

        # TODO 4/25 这边也是需要改进的一个地方，关于这个output有两个输出
        # 这边是不是需要计算两个准确度了
        _, avg_acc, cnt, pred = accuracy(output_h.detach().cpu().numpy(),
                                         target.detach().cpu().numpy())





def validate(config, val_loader, val_dataset, model, criterions: list, output_dir,
             tb_log_dir, writer_dict=None):
    pass


class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count if self.count != 0 else 0